from typing import List

from backend.app.common.jwt import CurrentUser
from backend.app.database.db_es import es_conn
from backend.app.schemas.device_v2 import DeviceSchemaV2
from backend.app.utils.const import get_index_by_platform


class TenderDeviceMatchChain:
    """【Pinecone向量数据库】按照最大相关性查询"""
    _top_k = 3  # 需要查处几个相似

    def __init__(self, platform: str, shop_id: int, top_k: int = 3):
        self._index_name = self.get_index_by_platform(platform)
        self._rdb = es_conn.get_langchain_client(self._index_name)
        self._sid = shop_id
        self._top_k = top_k

    @staticmethod
    def get_index_by_platform(claim: CurrentUser):
        """根据用户信息获取索引名称"""
        index_name = get_index_by_platform(claim.platform)
        return index_name

    def get_rdb(self):
        return self._rdb

    async def base_search(self, query: str, **kwargs):
        rdb = self.get_rdb()

        condition = dict()
        if self._sid > 0:
            condition["sid"] = self._sid
        result = await rdb.amax_marginal_relevance_search(
            # result = await rdb.asimilarity_search(
            query=query,
            k=self._top_k,
            filter=condition,
            fetch_k=50,
            lambda_mult=0.8,
            namespace=kwargs.get("namespace")
        )
        return result

    async def search(self, query: str, **kwargs):
        result = await self.base_search(query, **kwargs)
        return self.format_result(result)

    async def search_list(self, query: str, **kwargs):
        """搜索设备列表"""
        result = await self.base_search(query, **kwargs)
        return self.format_to_list(result)

    @staticmethod
    def format_to_list(score_docs: List) -> List[DeviceSchemaV2]:
        """将搜索结果转为设备列表"""
        lists = []
        if not score_docs:
            return lists
        for data in score_docs:
            item = DeviceSchemaV2(**data.metadata)
            lists.append(item)
        return lists

    @staticmethod
    def format_result(score_docs: List) -> DeviceSchemaV2:
        device = DeviceSchemaV2()
        if not score_docs:
            return device
        child_list = []
        for data in score_docs:
            item = DeviceSchemaV2(**data.metadata)
            if device.gid < 1:
                device = item
            else:
                child_list.append(item)
        if device and device.gid > 0:
            device.child_list = child_list
        return device


class TenderDeviceMatchScoreChain:
    """【ES向量数据库】 按照评分查询"""
    _top_k = 3  # 需要查处几个相似

    def __init__(self, platform: str, top_k: int = 3):
        self._index_name = self.get_index_by_platform(platform)
        self._rdb = es_conn.get_langchain_client(self._index_name)
        self._top_k = top_k

    @staticmethod
    def get_index_by_platform(platform: str):
        """根据用户信息获取索引名称"""
        index_name = get_index_by_platform(platform)
        return index_name

    def get_rdb(self):
        return self._rdb

    async def base_search(self, query: str, **kwargs):
        rdb = self.get_rdb()
        condition = list()
        if kwargs.get("country"):
            item = {"match": {"metadata.country": kwargs.get("country")}}
            condition.append(item)
        if kwargs.get("country_title"):
            item = {"match": {"metadata.country_title": kwargs.get("country_title")}}
            condition.append(item)
        result = await rdb.asimilarity_search_with_score(
            query=query,
            k=self._top_k,
            condition=condition,
        )
        return result

    async def search(self, query: str, **kwargs):
        """设备列表搜索-->返回主要设备和相似设备"""
        result = await self.base_search(query, **kwargs)
        return self.format_result(result)

    async def search_list(self, query: str, **kwargs):
        """搜索设备列表"""
        result = await self.base_search(query, **kwargs)
        return self.format_to_list(result)

    @staticmethod
    def format_to_list(score_docs: List) -> List[DeviceSchemaV2]:
        """将搜索结果转为设备列表"""
        lists = []
        if not score_docs:
            return lists
        for data, score in score_docs:
            item = DeviceSchemaV2(**data.metadata)
            lists.append(item)
        return lists

    @staticmethod
    def format_result(score_docs: List) -> DeviceSchemaV2:
        """列出结果中评分最高的设备，其他设备列表子设备列表"""
        device = DeviceSchemaV2()
        if not score_docs:
            return device
        for data, score in score_docs:
            if score < 0.3:
                continue
            item = DeviceSchemaV2(**data.metadata)
            if device.gid < 1:
                device = item
            else:
                device.child_list.append(item)
        return device
